{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this cell to install dependencies\n",
    "%pip install sf-hamilton[visualization]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hamilton UI Adapter [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dagworks-inc/hamilton/blob/main/examples/ray/ray_Hamilton_UI_tracking/hamilton_notebook.ipynb) [![GitHub badge](https://img.shields.io/badge/github-view_source-2b3137?logo=github)](https://github.com/apache/hamilton/blob/main/examples/ray/ray_Hamilton_UI_tracking/hamilton_notebook.ipynb)\n",
    "\n",
    "\n",
    "Needs a running instance of Hamilton UI: https://hamilton.apache.org/concepts/ui/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hamilton_sdk.adapters import HamiltonTracker\n",
    "\n",
    "# Inputs required to track into correct project in the UI\n",
    "project_id = 2\n",
    "username = \"admin\"\n",
    "\n",
    "tracker_ray = HamiltonTracker(\n",
    "        project_id=project_id,\n",
    "        username=username,\n",
    "        dag_name=\"telemetry_with_ray\",)\n",
    "\n",
    "tracker_without_ray = HamiltonTracker(\n",
    "        project_id=project_id,\n",
    "        username=username,\n",
    "        dag_name=\"telemetry_without_ray\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ray adapter\n",
    "\n",
    "https://hamilton.apache.org/reference/graph-adapters/RayGraphAdapter/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hamilton import base\n",
    "from hamilton.plugins.h_ray import RayGraphAdapter\n",
    "\n",
    "rga = RayGraphAdapter(result_builder=base.PandasDataFrameResult())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Importing Hamilton and the DAG modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hamilton import driver\n",
    "import ray_lineage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    dr_ray = driver.Builder().with_modules(ray_lineage).with_adapters(rga, tracker_ray).build()\n",
    "    result_ray = dr_ray.execute(\n",
    "        final_vars=[\n",
    "            \"node_5s\",\n",
    "            \"node_1s_error\",\n",
    "            \"add_1_to_previous\",\n",
    "        ]\n",
    "    )\n",
    "    print(result_ray)\n",
    "\n",
    "except ValueError:\n",
    "    print(\"UI should display failure.\")\n",
    "finally:\n",
    "    dr_without_ray = driver.Builder().with_modules(ray_lineage).with_adapters(tracker).build()\n",
    "    result_without_ray = dr_without_ray.execute(final_vars=[\"node_5s\", \"add_1_to_previous\"])\n",
    "    print(result_without_ray)  \n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
